import random
from IPython.core.display import SVG
import pyomo.environ as pyo
from pysat.solvers import Solver
from pysat.formula import CNF 
import py_svg_combinatorics as psc
from ipywidgets import widgets, HBox
from collections import Counter
from pprint import pprint
from random import randint
import numpy as np
from IPython.display import IFrame
import IPython
from copy import copy
import os
from pathlib import Path
nbname = ''
try:
    nbname = __vsc_ipynb_file__
except:
    if 'COCALC_JUPYTER_FILENAME' in os.environ:
        nbname = os.environ['COCALC_JUPYTER_FILENAME']    
title_ = Path(nbname).stem.replace('-', '_').title()    
IFrame(f'https://discopal.ispras.ru/index.php?title=Hardprob/{title_}&useskin=cleanmonobook', width=1280, height=300)

phys_regs_amount = 3
phys_regs_costs = np.arange(1, phys_regs_amount + 1)

basic_block_size = 16
instruction_uses_max = int(phys_regs_amount / 2)

virtual_alives = [0, 1]
virtual_uses = []

for _ in range(basic_block_size):
    # Каждая инструкция использует до instruction_uses_max виртуальных регистров
    virtual_uses.append([f"{i:02}" for i in np.random.choice(virtual_alives, instruction_uses_max)])
    
    # Инструкция с шансом 0.5 может определить новый виртуальный регистр
    if np.random.choice([True, False]):
        virtual_alives.append(virtual_alives[-1] + 1)

svg = psc.subsets2svg(virtual_uses)
SVG(data=svg)

def get_model(instruction_uses, phys_regs_costs, virtual_regs_amount):
    model = pyo.ConcreteModel()
    
    # Множество физических регистров
    model.phys_regs = pyo.Set(initialize=range(len(phys_regs_costs)))
    
    # Сопоставление каждому физическому регистру его стоимости
    model.phys_costs = pyo.Param(model.phys_regs, initialize=phys_regs_costs)
    
    # Множество виртуальных регистров
    model.virt_regs = pyo.Set(initialize=range(virtual_regs_amount))
    
    # Хотим составить отображение из виртуальных регистров на физические
    model.alloca = pyo.Var(model.virt_regs, model.phys_regs, domain=pyo.Binary)
    
    # Составляем множества плохих пар виртуальных регистров (встречаемых в используемых для одинаковых инструкций)
    model.bad_pairs = pyo.Set(initialize=sorted(set([(int(i), int(j)) for uses in instruction_uses for i in uses for j in uses if i < j])))
    
    # ЦЕЛЬ: минимизировать стоимость использования физических регистров после аллокации
    model.obj = pyo.Objective(expr=sum(model.phys_costs[p] * model.alloca[v, p] 
                                for v in model.virt_regs 
                                for p in model.phys_regs))
    
    # ОГРАНИЧЕНИЕ: каждому виртуальному регистру сопоставляем ровно 1 физический
    @model.Constraint(model.virt_regs)
    def виртуальному_один_физический(m, v):
        return sum(model.alloca[v, phys] for phys in m.phys_regs) == 1
    
    # ОГРАНИЧЕНИЕ: каждый физический регистр должен использоваться каждой инструкцией не более одного раза
    @model.Constraint(model.bad_pairs, model.phys_regs)
    def использование_одного_физического(m, first, second, reg):
        return (m.alloca[first, reg] + m.alloca[second, reg]) <= 1
    
    return model
    

m = get_model(virtual_uses, phys_regs_costs, len(virtual_alives))
m.pprint()

solver = pyo.SolverFactory('cbc')
solver.solve(m).write()

for v in m.component_data_objects(pyo.Var):
    if v.value and v.value > 0:
        print(str(v), v.value)

phys_uses = []
for use in virtual_uses:
    phys_uses.append([f"{p:02}" for v in use for p in range(phys_regs_amount) if pyo.value(m.alloca[int(v), p])])
svg = psc.subsets2svg(phys_uses)
SVG(data=svg)

sat_vars_amount = 5
sat_clauses_amount = 20

cnf3 = CNF(from_clauses=psc.rand3cnf(sat_clauses_amount, sat_vars_amount))
cnf3.clauses

solver = Solver(bootstrap_with=cnf3)
res = solver.solve()
res

NEUTRAL = 0
TRUE = 1
FALSE = 2

vertexes_amount = 3
edges = [(NEUTRAL, TRUE), (TRUE, FALSE), (FALSE, NEUTRAL)]

def create_var():
    global vertexes_amount
    
    new_var = vertexes_amount
    vertexes_amount += 1

    return new_var

def ident_var():
    global edges

    new_var = create_var()
    edges.append((new_var, NEUTRAL))

    return new_var

def negate_var(to_negate):
    global edges
    
    negative_var = ident_var()
    edges.append((negative_var, to_negate))

    return negative_var

def or_var(lhs, rhs):
    global edges
    global vertexes_amount

    res_var = ident_var()

    tmps = list(range(vertexes_amount, vertexes_amount + 4))
    vertexes_amount += 4

    edges.append((TRUE, tmps[0]))

    edges.append((tmps[0], tmps[1]))
    edges.append((tmps[0], res_var))

    edges.append((tmps[1], lhs))
    edges.append((tmps[1], rhs))

    edges.append((lhs, tmps[2]))
    edges.append((rhs, tmps[3]))

    edges.append((tmps[2], tmps[3]))
    edges.append((tmps[2], res_var))
    edges.append((tmps[3], res_var))

    return res_var

import functools

var_converter = {}
for sat_var in range(1, sat_vars_amount + 1):
    var_converter[sat_var] = ident_var()

negate_converter = {}
for sat_var in range(1, sat_vars_amount + 1):
    negate_converter[sat_var] = negate_var(var_converter[sat_var])

converted_clauses = []
for clause in cnf3.clauses:
    converted_clauses.append(functools.reduce(or_var, map(lambda x: var_converter[x] if x > 0 else negate_converter[-x], clause)))

for clause in converted_clauses:
    edges.append((clause, NEUTRAL))
    edges.append((clause, FALSE))

m = get_model(edges, phys_regs_costs, vertexes_amount)
m.pprint()

cbc_solver = pyo.SolverFactory('cbc')
cbc_solver.solve(m).write()

false_reg = [p for p in range(phys_regs_amount) if pyo.value(m.alloca[FALSE, p])]
true_reg = [p for p in range(phys_regs_amount) if pyo.value(m.alloca[TRUE, p])]

false_reg = false_reg[0]
true_reg = true_reg[0]

answer = []
for sat_var in range(1, sat_vars_amount + 1):
    is_false = pyo.value(m.alloca[var_converter[sat_var], false_reg])
    is_true = pyo.value(m.alloca[var_converter[sat_var], true_reg])
    
    assert(is_false or is_true)
    if is_true:
        answer.append(sat_var)
    else:
        answer.append(-sat_var)

print(answer)
print(list(solver.enum_models()))